#!/usr/bin/env python3
"""
SQL 查询执行引擎

本模块负责安全地执行 SQL 查询并处理结果。
主要功能包括：
- SQL 查询验证和清理
- 查询执行和结果处理
- 查询性能监控
- 安全检查和权限控制
- 结果格式化和分页
"""

import re
import time
from typing import Any, Dict, List, Optional, Tuple, Union, Set
from dataclasses import dataclass
from datetime import datetime
from enum import Enum

from .logger import get_logger, performance_logger, log_execution_time
from .connection_pool import DatabaseConnectionPool


class QueryValidationError(Exception):
    """
    查询验证错误异常
    
    当 SQL 查询验证失败时抛出此异常。
    """
    
    def __init__(self, message: str, query: Optional[str] = None):
        """
        初始化查询验证错误
        
        Args:
            message: 错误消息
            query: 导致错误的查询语句
        """
        self.message = message
        self.query = query
        super().__init__(message)


class QueryExecutionError(Exception):
    """
    查询执行错误异常
    
    当 SQL 查询执行失败时抛出此异常。
    """
    
    def __init__(self, message: str, query: Optional[str] = None, original_error: Optional[Exception] = None):
        """
        初始化查询执行错误
        
        Args:
            message: 错误消息
            query: 导致错误的查询语句
            original_error: 原始异常
        """
        self.message = message
        self.query = query
        self.original_error = original_error
        super().__init__(message)


class QuerySecurityError(Exception):
    """
    查询安全错误异常
    
    当 SQL 查询存在安全风险时抛出此异常。
    """
    
    def __init__(self, message: str, query: Optional[str] = None, security_issue: Optional[str] = None):
        """
        初始化查询安全错误
        
        Args:
            message: 错误消息
            query: 导致错误的查询语句
            security_issue: 安全问题描述
        """
        self.message = message
        self.query = query
        self.security_issue = security_issue
        super().__init__(message)


class QueryType(Enum):
    """
    SQL 查询类型枚举
    """
    SELECT = "SELECT"
    INSERT = "INSERT"
    UPDATE = "UPDATE"
    DELETE = "DELETE"
    CREATE = "CREATE"
    DROP = "DROP"
    ALTER = "ALTER"
    SHOW = "SHOW"
    DESCRIBE = "DESCRIBE"
    EXPLAIN = "EXPLAIN"
    UNKNOWN = "UNKNOWN"


@dataclass
class QueryResult:
    """
    查询结果数据结构
    """
    success: bool
    data: Optional[List[Dict[str, Any]]] = None
    affected_rows: Optional[int] = None
    execution_time: Optional[float] = None
    query_type: Optional[QueryType] = None
    error: Optional[str] = None
    warning: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class QueryValidationResult:
    """
    查询验证结果
    """
    is_valid: bool
    query_type: QueryType
    is_safe: bool
    warnings: List[str]
    errors: List[str]
    cleaned_query: str


class SQLQueryEngine:
    """
    SQL 查询执行引擎
    
    提供安全的 SQL 查询执行功能，包括查询验证、执行监控和结果处理。
    """
    
    # 危险的 SQL 关键词模式
    DANGEROUS_PATTERNS = [
        r'\b(DROP|TRUNCATE|DELETE)\s+(?!.*WHERE)',  # 没有 WHERE 条件的危险操作
        r'\bALTER\s+USER',  # 用户权限修改
        r'\bGRANT\s+',  # 权限授予
        r'\bREVOKE\s+',  # 权限撤销
        r'\bCREATE\s+USER',  # 创建用户
        r'\bSET\s+PASSWORD',  # 设置密码
        r'\bLOAD\s+DATA',  # 数据加载
        r'\bSELECT\s+.*INTO\s+OUTFILE',  # 文件输出
        r'\bUNION\s+.*SELECT.*--',  # SQL 注入模式
        r'/\*.*\*/',  # 注释可能隐藏恶意代码
        r'--;',  # 注释可能隐藏恶意代码
    ]
    
    # 只读查询类型
    READ_ONLY_TYPES = {
        QueryType.SELECT,
        QueryType.SHOW,
        QueryType.DESCRIBE,
        QueryType.EXPLAIN
    }
    
    # 写操作查询类型
    WRITE_TYPES = {
        QueryType.INSERT,
        QueryType.UPDATE,
        QueryType.DELETE,
        QueryType.CREATE,
        QueryType.DROP,
        QueryType.ALTER
    }
    
    def __init__(self, db_pool: DatabaseConnectionPool, 
                 enable_write_operations: bool = False,
                 max_result_size: int = 10000,
                 query_timeout: int = 30):
        """
        初始化查询执行引擎
        
        Args:
            db_pool: 数据库连接池
            enable_write_operations: 是否允许写操作
            max_result_size: 最大结果集大小
            query_timeout: 查询超时时间（秒）
        """
        self.logger = get_logger(__name__)
        self.db_pool = db_pool
        self.enable_write_operations = enable_write_operations
        self.max_result_size = max_result_size
        self.query_timeout = query_timeout
        
        # 编译正则表达式模式
        self.dangerous_patterns = [re.compile(pattern, re.IGNORECASE) 
                                 for pattern in self.DANGEROUS_PATTERNS]
        
        # 查询统计
        self.query_stats = {
            "total_queries": 0,
            "successful_queries": 0,
            "failed_queries": 0,
            "read_queries": 0,
            "write_queries": 0,
            "blocked_queries": 0
        }
        
        # 为了向后兼容，提供stats属性
        self.stats = self.query_stats
    
    @log_execution_time()
    async def execute_query(self, query: str, 
                          parameters: Optional[Dict[str, Any]] = None,
                          limit: Optional[int] = None,
                          validate_only: bool = False) -> QueryResult:
        """
        执行 SQL 查询
        
        Args:
            query: SQL 查询语句
            parameters: 查询参数（用于参数化查询）
            limit: 结果限制数量
            validate_only: 是否仅验证查询而不执行
            
        Returns:
            QueryResult: 查询结果
        """
        start_time = time.time()
        self.query_stats["total_queries"] += 1
        
        try:
            # 验证查询
            validation_result = self._validate_query(query)
            
            if not validation_result.is_valid:
                self.query_stats["blocked_queries"] += 1
                return QueryResult(
                    success=False,
                    error=f"查询验证失败: {'; '.join(validation_result.errors)}",
                    query_type=validation_result.query_type,
                    execution_time=time.time() - start_time
                )
            
            if not validation_result.is_safe:
                self.query_stats["blocked_queries"] += 1
                return QueryResult(
                    success=False,
                    error=f"查询被安全检查阻止: {'; '.join(validation_result.warnings)}",
                    query_type=validation_result.query_type,
                    execution_time=time.time() - start_time
                )
            
            # 检查写操作权限
            if (validation_result.query_type in self.WRITE_TYPES and 
                not self.enable_write_operations):
                self.query_stats["blocked_queries"] += 1
                return QueryResult(
                    success=False,
                    error="写操作被禁用",
                    query_type=validation_result.query_type,
                    execution_time=time.time() - start_time
                )
            
            # 如果只是验证，返回成功结果
            if validate_only:
                return QueryResult(
                    success=True,
                    query_type=validation_result.query_type,
                    execution_time=time.time() - start_time,
                    metadata={"validation_warnings": validation_result.warnings}
                )
            
            # 执行查询
            result = await self._execute_validated_query(
                validation_result.cleaned_query,
                validation_result.query_type,
                parameters,
                limit
            )
            
            # 更新统计
            if result.success:
                self.query_stats["successful_queries"] += 1
                if validation_result.query_type in self.READ_ONLY_TYPES:
                    self.query_stats["read_queries"] += 1
                else:
                    self.query_stats["write_queries"] += 1
            else:
                self.query_stats["failed_queries"] += 1
            
            # 记录性能数据
            performance_logger.log_query_performance(
                query=validation_result.cleaned_query,
                execution_time=result.execution_time or 0,
                result_count=len(result.data) if result.data else 0,
                query_type=validation_result.query_type.value,
                success=result.success
            )
            
            return result
            
        except Exception as e:
            self.query_stats["failed_queries"] += 1
            self.logger.error("执行查询时发生未知错误", query=query, error=str(e))
            
            return QueryResult(
                success=False,
                error=f"执行查询时发生错误: {str(e)}",
                execution_time=time.time() - start_time
            )
    
    def _validate_query(self, query: str) -> QueryValidationResult:
        """
        验证 SQL 查询
        
        Args:
            query: SQL 查询语句
            
        Returns:
            QueryValidationResult: 验证结果
        """
        errors = []
        warnings = []
        
        # 基本检查
        if not query or not query.strip():
            errors.append("查询语句不能为空")
            return QueryValidationResult(
                is_valid=False,
                query_type=QueryType.UNKNOWN,
                is_safe=False,
                warnings=warnings,
                errors=errors,
                cleaned_query=query
            )
        
        # 清理查询语句
        cleaned_query = self._clean_query(query)
        
        # 检测查询类型
        query_type = self._detect_query_type(cleaned_query)
        
        # 长度检查
        if len(cleaned_query) > 50000:  # 50KB 限制
            errors.append("查询语句过长")
        
        # 安全检查
        is_safe = self._check_query_safety(cleaned_query, warnings)
        
        # 语法基本检查
        syntax_valid = self._basic_syntax_check(cleaned_query, query_type, errors)
        
        return QueryValidationResult(
            is_valid=len(errors) == 0 and syntax_valid,
            query_type=query_type,
            is_safe=is_safe,
            warnings=warnings,
            errors=errors,
            cleaned_query=cleaned_query
        )
    
    def _clean_query(self, query: str) -> str:
        """
        清理查询语句
        
        Args:
            query: 原始查询语句
            
        Returns:
            str: 清理后的查询语句
        """
        # 移除多余的空白字符
        cleaned = re.sub(r'\s+', ' ', query.strip())
        
        # 移除行尾分号（如果存在）
        cleaned = cleaned.rstrip(';')
        
        return cleaned
    
    def _detect_query_type(self, query: str) -> QueryType:
        """
        检测查询类型
        
        Args:
            query: SQL 查询语句
            
        Returns:
            QueryType: 查询类型
        """
        query_upper = query.upper().strip()
        
        if query_upper.startswith('SELECT'):
            return QueryType.SELECT
        elif query_upper.startswith('INSERT'):
            return QueryType.INSERT
        elif query_upper.startswith('UPDATE'):
            return QueryType.UPDATE
        elif query_upper.startswith('DELETE'):
            return QueryType.DELETE
        elif query_upper.startswith('CREATE'):
            return QueryType.CREATE
        elif query_upper.startswith('DROP'):
            return QueryType.DROP
        elif query_upper.startswith('ALTER'):
            return QueryType.ALTER
        elif query_upper.startswith('SHOW'):
            return QueryType.SHOW
        elif query_upper.startswith('DESCRIBE') or query_upper.startswith('DESC'):
            return QueryType.DESCRIBE
        elif query_upper.startswith('EXPLAIN'):
            return QueryType.EXPLAIN
        else:
            return QueryType.UNKNOWN
    
    def _check_query_safety(self, query: str, warnings: List[str]) -> bool:
        """
        检查查询安全性
        
        Args:
            query: SQL 查询语句
            warnings: 警告列表（会被修改）
            
        Returns:
            bool: 是否安全
        """
        is_safe = True
        
        # 检查危险模式
        for pattern in self.dangerous_patterns:
            if pattern.search(query):
                warnings.append(f"检测到潜在危险操作: {pattern.pattern}")
                is_safe = False
        
        # 检查多语句
        if ';' in query.rstrip(';'):
            warnings.append("不允许执行多条语句")
            is_safe = False
        
        # 检查嵌套查询深度
        nested_level = query.upper().count('SELECT')
        if nested_level > 5:
            warnings.append("嵌套查询层级过深")
            is_safe = False
        
        return is_safe
    
    def _basic_syntax_check(self, query: str, query_type: QueryType, 
                          errors: List[str]) -> bool:
        """
        基本语法检查
        
        Args:
            query: SQL 查询语句
            query_type: 查询类型
            errors: 错误列表（会被修改）
            
        Returns:
            bool: 语法是否有效
        """
        query_upper = query.upper()
        
        # 检查括号匹配
        if query.count('(') != query.count(')'):
            errors.append("括号不匹配")
            return False
        
        # 检查引号匹配
        single_quotes = query.count("'") - query.count("\\'")
        double_quotes = query.count('"') - query.count('\\"')
        
        if single_quotes % 2 != 0:
            errors.append("单引号不匹配")
            return False
        
        if double_quotes % 2 != 0:
            errors.append("双引号不匹配")
            return False
        
        # 特定查询类型的语法检查
        if query_type == QueryType.SELECT:
            if 'FROM' not in query_upper and 'DUAL' not in query_upper:
                errors.append("SELECT 查询缺少 FROM 子句")
                return False
        
        elif query_type == QueryType.INSERT:
            if 'INTO' not in query_upper:
                errors.append("INSERT 查询缺少 INTO 子句")
                return False
        
        elif query_type == QueryType.UPDATE:
            if 'SET' not in query_upper:
                errors.append("UPDATE 查询缺少 SET 子句")
                return False
        
        return True
    
    async def _execute_validated_query(self, query: str, query_type: QueryType,
                                     parameters: Optional[Dict[str, Any]] = None,
                                     limit: Optional[int] = None) -> QueryResult:
        """
        执行已验证的查询
        
        Args:
            query: 已验证的 SQL 查询语句
            query_type: 查询类型
            parameters: 查询参数
            limit: 结果限制
            
        Returns:
            QueryResult: 查询结果
        """
        start_time = time.time()
        
        try:
            # 应用结果限制
            if query_type in self.READ_ONLY_TYPES and limit:
                effective_limit = min(limit, self.max_result_size)
                if 'LIMIT' not in query.upper():
                    query += f" LIMIT {effective_limit}"
            
            # 执行查询
            if query_type in self.READ_ONLY_TYPES:
                # 读操作
                data = self.db_pool.execute_query(
                    query,
                    parameters
                )
                
                return QueryResult(
                    success=True,
                    data=data,
                    query_type=query_type,
                    execution_time=time.time() - start_time,
                    metadata={
                        "result_count": len(data) if data else 0,
                        "truncated": len(data) >= self.max_result_size if data else False
                    }
                )
            
            else:
                # 写操作
                result = self.db_pool.execute_query(
                    query,
                    parameters
                )
                
                # 对于写操作，返回受影响的行数
                affected_rows = 0
                if isinstance(result, list) and len(result) > 0:
                    # 如果返回的是字典列表，尝试获取affected_rows
                    if isinstance(result[0], dict) and 'affected_rows' in result[0]:
                        affected_rows = result[0]['affected_rows']
                elif isinstance(result, int):
                    affected_rows = result
                
                return QueryResult(
                    success=True,
                    affected_rows=affected_rows,
                    query_type=query_type,
                    execution_time=time.time() - start_time
                )
        
        except Exception as e:
            self.logger.error("执行查询时发生数据库错误", query=query, error=str(e))
            
            return QueryResult(
                success=False,
                error=str(e),
                query_type=query_type,
                execution_time=time.time() - start_time
            )
    
    def get_query_statistics(self) -> Dict[str, Any]:
        """
        获取查询统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        return {
            **self.query_stats,
            "success_rate": (
                self.query_stats["successful_queries"] / 
                max(self.query_stats["total_queries"], 1)
            ) * 100,
            "timestamp": datetime.now().isoformat()
        }
    
    def reset_statistics(self) -> None:
        """
        重置查询统计信息
        """
        for key in self.query_stats:
            self.query_stats[key] = 0
        
        self.logger.info("查询统计信息已重置")
    
    async def explain_query(self, query: str) -> QueryResult:
        """
        分析查询执行计划
        
        Args:
            query: SQL 查询语句
            
        Returns:
            QueryResult: 执行计划结果
        """
        if not query.upper().strip().startswith('SELECT'):
            return QueryResult(
                success=False,
                error="只能分析 SELECT 查询的执行计划"
            )
        
        explain_query = f"EXPLAIN {query}"
        return await self.execute_query(explain_query, validate_only=False)
    
    async def get_table_info(self, table_name: str) -> QueryResult:
        """
        获取表信息
        
        Args:
            table_name: 表名
            
        Returns:
            QueryResult: 表信息
        """
        # 验证表名（防止 SQL 注入）
        if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', table_name):
            return QueryResult(
                success=False,
                error="无效的表名"
            )
        
        describe_query = f"DESCRIBE {table_name}"
        return await self.execute_query(describe_query, validate_only=False)
    
    def sanitize_sql(self, query: str) -> str:
        """
        清理和净化SQL查询语句
        
        Args:
            query: 原始SQL查询语句
            
        Returns:
            str: 清理后的SQL查询语句
        """
        if not query:
            return ""
        
        # 移除前后空白
        cleaned = query.strip()
        
        # 移除末尾的分号
        while cleaned.endswith(';'):
            cleaned = cleaned[:-1].strip()
        
        # 移除注释
        if '--' in cleaned:
            cleaned = cleaned.split('--')[0].strip()
        
        # 规范化空白字符
        import re
        cleaned = re.sub(r'\s+', ' ', cleaned)
        
        return cleaned
    
    def get_query_type(self, query: str) -> str:
        """
        获取查询类型
        
        Args:
            query: SQL查询语句
            
        Returns:
            str: 查询类型
        """
        query_type = self._detect_query_type(query)
        return query_type.value
    
    def is_read_only_query(self, query: str) -> bool:
        """
        检查查询是否为只读查询
        
        Args:
            query: SQL查询语句
            
        Returns:
            bool: 是否为只读查询
        """
        query_type = self._detect_query_type(query)
        return query_type in self.READ_ONLY_TYPES
    
    async def execute_batch_queries(self, queries: List[Union[str, Tuple[str, Any]]]) -> List[QueryResult]:
        """
        批量执行SQL查询
        
        Args:
            queries: SQL查询语句列表，可以是字符串或(query, parameters)元组
            
        Returns:
            List[QueryResult]: 查询结果列表
        """
        results = []
        
        for query_item in queries:
            if isinstance(query_item, tuple):
                query, query_params = query_item
            else:
                query = query_item
                query_params = None
                
            result = await self.execute_query(query, query_params)
            results.append(result)
            
            # 如果查询失败且不是只读查询，停止执行后续查询
            if not result.success and not self.is_read_only_query(query):
                break
        
        return results
    
    def validate_sql(self, query: str) -> None:
        """验证SQL查询"""
        validation_result = self._validate_query(query)
        
        if not validation_result.is_valid:
            raise QueryValidationError(f"查询验证失败: {'; '.join(validation_result.errors)}")
        
        if not validation_result.is_safe:
            raise QuerySecurityError(f"查询被安全检查阻止: {'; '.join(validation_result.warnings)}")
        
        # 检查写操作权限
        if (validation_result.query_type in self.WRITE_TYPES and 
            not self.enable_write_operations):
            raise QuerySecurityError("写操作被禁用")
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = self.stats.copy()
        
        # 计算成功率
        total = stats.get('total_queries', 0)
        successful = stats.get('successful_queries', 0)
        if total > 0:
            stats['success_rate'] = (successful / total) * 100.0
        else:
            stats['success_rate'] = 0.0
        
        # 添加配置信息
        stats['enable_write_operations'] = self.enable_write_operations
        stats['max_result_size'] = self.max_result_size
        stats['query_timeout'] = self.query_timeout
        
        return stats
    
    def reset_stats(self) -> None:
        """重置统计信息"""
        self.stats = {
            'total_queries': 0,
            'successful_queries': 0,
            'failed_queries': 0,
            'total_execution_time': 0.0,
            'avg_execution_time': 0.0
        }


def create_query_engine(db_pool: DatabaseConnectionPool, 
                       enable_write_operations: bool = False,
                       max_result_size: int = 10000,
                       query_timeout: int = 30) -> SQLQueryEngine:
    """
    创建查询执行引擎
    
    Args:
        db_pool: 数据库连接池
        enable_write_operations: 是否允许写操作
        max_result_size: 最大结果集大小
        query_timeout: 查询超时时间
        
    Returns:
        SQLQueryEngine: 查询执行引擎实例
    """
    return SQLQueryEngine(
        db_pool=db_pool,
        enable_write_operations=enable_write_operations,
        max_result_size=max_result_size,
        query_timeout=query_timeout
    )


# 为了向后兼容，提供 QueryEngine 别名
QueryEngine = SQLQueryEngine